import torch

from defense.base_defense import BaseDefense

class Defense(BaseDefense):
    def __init__(self, model_name, config):
        super().__init__('RC', model_name, config, retrain = False, ignore_output = True)
        if self.config['search']:
            self.radius = self.search_for_radius()
        else:
            self.radius = round(self.config['radius'], 2)
        print('Radius is set to: {}'.format(self.radius))

    def search_for_radius(self):
        train_loader, valid_loader = self.model.get_train_valid_loader(batch_size = 1000,
                                                                        valid_size = 0.02,
                                                                        transform_train = True,
                                                                        shuffle = True)
        self.model.load_raw_state()

        self.model.eval()
        with torch.no_grad():
            # compute the original classification accuracy on validation dataset
            val_acc = self.model.test(valid_loader)
            print('<--- original classification accuracy on validation dataset is {:.4f} --->'.format(val_acc))

            # learn the radius through a search process
            total_step = int((self.config['radius_max'] - self.config['radius_min']) / self.config['radius_step'])
            for index in range(total_step):

                # update the radius
                self.radius = self.config['radius_min'] + self.config['radius_step'] * (index + 1)
                print('Searching radius = {}'.format(self.radius))

                # calculate the accuracy of region-based classification on validation dataset
                rc_acc = self.test(valid_loader)

                if (val_acc - rc_acc) > 1e-2:
                    return round(self.radius - self.config['radius_step'], 2)

            return self.config['radius_max']

    def region_based_classification_single(self, sample):
        with torch.no_grad():
            # prepare the hypercube samples (size=num_points) for the sample (size=1)
            hypercube_samples = sample.repeat(self.config['num_points'], 1, 1, 1)
            random_space = torch.Tensor(*hypercube_samples.size()).to(self.model.device).float()
            random_space.uniform_(-self.radius, self.radius)
            hypercube_samples = torch.clamp(hypercube_samples + random_space, min=0.0, max=1.0)

            # predicting for hypercube samples
            hypercube_preds = self.model(hypercube_samples)
            hypercube_labels = torch.argmax(hypercube_preds, dim=1)

            # voting for predicted labels
            bin_count = torch.bincount(hypercube_labels)
            rc_label = torch.argmax(bin_count, dim=0)

            out = torch.zeros(hypercube_preds.size(1)).to(self.model.device)
            out[rc_label] = 1
            return out

    def predict(self, samples):
        self.model.eval()
        outputs = []
        for i in range(samples.size(0)):
            x = samples[i: i + 1]
            output = self.region_based_classification_single(x)
            outputs.append(output)
        return torch.stack(outputs)
